import json
import math
import re
import sys
import time

from bins.tool import interpolate_trajectory_with_heading
from bins.vehicleParams import VehicleParams
import numpy as np
from tool.visual.popWindows.pop import *
from PyQt5 import QtCore, QtGui, QtWidgets
from tool.visual.VisualTuning import Ui_MainWindow
from PyQt5.Qt import *
from src.common.processedMatrix import ProcessedMatrix
from src.common.processedmeasurement import ProcessedMeasurement
import pyqtgraph as pg


class Visual(QMainWindow, Ui_MainWindow):
    def __init__(self):
        super(Visual, self).__init__()
        self.predict_traj = None
        self.cost = None
        self.setupUi(self)
        self.pen = None
        self.win = None
        self.vertical_layout = None
        self.dtype = np.dtype([
            ('x', np.float32),
            ('y', np.float32),
            ('h', np.float32),
            ('v', np.float32),
            ('t', np.float32)
        ])
        self.N = 25
        self.dt = 0.1
        self.trajectory_file = None
        self.current_state = None
        self.current_state_mpc = None
        self.current_control_steer = None
        self.current_control_thr = None
        self.trajectory_index = 0
        self.trajectory_show = None
        self.max_index = None
        self.mark_slider = 0
        self.vehicle_conf_file = "../.././conf/vehicle_params.json"
        self.cost_conf_file = "../.././conf/cost_config.json"
        self.vehicle_params = None
        self.processed_measurement = ProcessedMeasurement()
        self.current_predict_traj = None
        self.current_predict_traj_show = None
        self.lateral_windows = None
        self.longitudinal_windows = None
        self.longitudinal = ProcessedMatrix()
        self.transverse = ProcessedMatrix()
        self.connect_bottom()
        self.InitMap()
        self.load_cost()
        self.init_cost()

    def connect_bottom(self):
        self.pushButton.clicked.connect(self.load_trajectory)
        self.pushButton_2.clicked.connect(self.load_vehicle)
        self.horizontalSlider.valueChanged.connect(self.exchange_index)
        self.horizontalSlider.sliderPressed.connect(lambda: self.change_mark(True))
        self.horizontalSlider.sliderReleased.connect(lambda: self.change_mark(False))
        self.pushButton_3.clicked.connect(self.init_windows_lateral)
        self.pushButton_4.clicked.connect(self.init_windows_longitudinal)
        self.lineEdit_3.returnPressed.connect(lambda: self.getResult(3))
        self.lineEdit_4.returnPressed.connect(lambda: self.getResult(3))
        self.lineEdit_5.returnPressed.connect(lambda: self.getResult(3))
        self.lineEdit_2.returnPressed.connect(lambda: self.getResult(3))
        self.lineEdit_6.returnPressed.connect(lambda: self.getResult(3))
        self.lineEdit_7.returnPressed.connect(lambda: self.getResult(3))
        self.lineEdit_8.returnPressed.connect(lambda: self.getResult(3))
        self.lineEdit_9.returnPressed.connect(lambda: self.getResult(3))
        pass

    def init_windows_lateral(self):
        if self.lateral_windows is None:
            self.lateral_windows = PopWindows("lateral")
            self.lateral_windows.init_windows("∆y", "∆h", "steer")
        self.lateral_windows.show()

    def init_windows_longitudinal(self):
        if self.longitudinal_windows is None:
            self.longitudinal_windows = PopWindows("longitudinal")
            self.longitudinal_windows.init_windows("s", "v", "a")
        self.longitudinal_windows.show()

    def change_mark(self, value: bool):
        self.mark_slider = value

    def getResult(self, mode):
        if self.vehicle_params is None:
            self.vehicle_params = VehicleParams()
            self.vehicle_params.load_vehicle_params()
            self.vehicle_params.calculate_axis()
            self.lineEdit_10.setText(self.vehicle_conf_file)

        t = self.horizontalSlider.value() / 100
        self.trajectory_index = int(t * self.max_index)
        self.current_state = list(self.processed_measurement.traj_point[self.trajectory_index])
        self.calculate_mpc_need_predict_traj()
        self.current_state_mpc = self.processed_measurement.calculate_current_state
        if mode == 1:
            self.calculate_transverse_mpc()
        elif mode == 2:
            self.calculate_longitudinal_mpc()
        elif mode == 3:
            self.calculate_transverse_mpc()
            self.calculate_longitudinal_mpc()
            self.generate_predict_traj()

    def exchange_index(self):
        if self.mark_slider:
            self.getResult(3)

    def update_current_predict_traj(self):
        if self.current_predict_traj_show is None:
            pass
        else:
            self.pen.removeItem(self.current_predict_traj_show)
        self.current_predict_traj_show = pg.PlotCurveItem(list(self.current_predict_traj['x']),
                                                          list(self.current_predict_traj['y']),
                                                          pen=pg.mkPen(color=(255, 63, 38), width=2, symbolBrush='r'))
        self.pen.addItem(self.current_predict_traj_show)

    def calculate_mpc_need_predict_traj(self):
        if self.current_state is None:
            return False
        self.processed_measurement.update_state(self.current_state)
        self.processed_measurement.calculate_ref_point()
        if self.processed_measurement.nearest_index + self.N < self.max_index:
            self.current_predict_traj = self.processed_measurement.traj_point[self.processed_measurement.nearest_index:
                                                                              self.processed_measurement.nearest_index + self.N]
        else:
            append_state = np.array([self.processed_measurement.traj_point[-1]] *
                                    (self.processed_measurement.nearest_index + self.N - self.max_index),
                                    dtype=self.processed_measurement.dtype)
            self.current_predict_traj = np.concatenate((
                self.processed_measurement.traj_point[self.processed_measurement.nearest_index:
                                                      self.max_index], append_state))

        self.update_current_predict_traj()

    def calculate_mpc_for_longitudinal(self):
        if self.current_state is None:
            return False

    def load_vehicle_config(self):
        self.lineEdit_10.setText(self.vehicle_conf_file)
        self.vehicle_params.load_vehicle_params(self.vehicle_conf_file)
        self.vehicle_params.calculate_axis()

    def load_vehicle(self):
        fname, ftype = QFileDialog.getOpenFileNames(self, "open File", '../../', 'json(*.json)')
        if fname:
            self.vehicle_conf_file = fname[0]
            self.load_vehicle_config()

    def load_cost(self):
        self.cost = json.loads(open(self.cost_conf_file, "r").read())

    def load_trajectory(self):
        fname, ftype = QFileDialog.getOpenFileNames(self, "open File", '../../', 'txt(*.txt)')
        if fname:
            self.trajectory_file = fname[0]
            self.lineEdit.setText(self.trajectory_file)
            try:
                trajectories = np.loadtxt(self.trajectory_file, delimiter=' ', dtype=self.dtype)
            except:
                trajectories = np.loadtxt(self.trajectory_file, delimiter=',', dtype=self.dtype)
            trajectories_unique = np.unique(trajectories, axis=0)
            # 提取轨迹中的变量（假设列顺序是 x, y, v, h）
            x = trajectories_unique['x']
            y = trajectories_unique['y']
            v = trajectories_unique['v']
            h = trajectories_unique['h'] # 航向角
            # 进行时间插值（目标时间步长 0.1s）
            t_, x_new, y_new, v_new, h_new = interpolate_trajectory_with_heading(x, y, v, h, dt=0.1)
            # 重新组合插值后的轨迹
            trajectories_processed = np.zeros(len(x_new), dtype=self.dtype)
            trajectories_processed['x'] = x_new
            trajectories_processed['y'] = y_new
            trajectories_processed['v'] = v_new
            trajectories_processed['h'] = h_new
            trajectories_processed['t'] = t_


            self.processed_measurement.predicted_step = self.N + 1
            self.processed_measurement.init_traj(trajectories_processed)
            self.max_index = len(trajectories_unique)
            self.pen.clear()
            self.show_trajectory()

    def show_trajectory(self):
        if self.processed_measurement.is_empty():
            return
        x_ = list(self.processed_measurement.traj_point['x'])
        y_ = list(self.processed_measurement.traj_point['y'])
        list_x = [x_[0]]
        list_y = [y_[0]]
        for i in range(1, len(x_)):
            if list_x[-1] == x_[i] and list_y[-1] == y_[i]:
                pass
            else:
                list_x.append(x_[i])
                list_y.append(y_[i])

        self.trajectory_show = pg.PlotCurveItem(list_x, list_y,
                                                pen=pg.mkPen(color=(0, 0, 100), width=2, symbolBrush='r'))
        self.pen.addItem(self.trajectory_show)

    def InitMap(self):
        self.vertical_layout = QVBoxLayout(self.widget_2)
        self.win = pg.GraphicsLayoutWidget(self.widget_2)
        self.win.setBackground((245, 245, 245))
        self.vertical_layout.addWidget(self.win)
        self.pen = self.win.addPlot()
        self.pen.setAspectLocked()
        s1 = pg.ScatterPlotItem()
        s2 = pg.ScatterPlotItem()
        self.pen.showGrid(x=True, y=True)
        self.pen.setLabel(axis="left", text="<span style='font-size:16px;color:black;font-family: Arial'>y(m)</span>")
        self.pen.setLabel(axis="bottom", text="<span style='font-size:16px;color:black;font-family: Arial'>x(m)</span>")
        myPen = pg.mkPen({'color': (0, 255, 0), 'width': 5})
        myPen2 = pg.mkPen({'color': (255, 0, 0), 'width': 5})
        self.pen.addItem(s1, pen=myPen, name="y2")
        self.pen.addItem(s2, pen=myPen2, name="y4")

    def init_cost(self):
        self.lineEdit_3.setText(str(self.cost["lateral"]["w1"]))
        self.lineEdit_4.setText(str(self.cost["lateral"]["w2"]))
        self.lineEdit_5.setText(str(self.cost["lateral"]["w3"]))
        self.lineEdit_2.setText(str(self.cost["lateral"]["w4"]))
        self.lineEdit_6.setText(str(self.cost["longitudinal"]["w1"]))
        self.lineEdit_7.setText(str(self.cost["longitudinal"]["w2"]))
        self.lineEdit_8.setText(str(self.cost["longitudinal"]["w3"]))
        self.lineEdit_9.setText(str(self.cost["longitudinal"]["w4"]))

    def get_transverse_cost(self):
        try:
            w1 = float(self.lineEdit_3.text())
            w2 = float(self.lineEdit_4.text())
            w3 = float(self.lineEdit_5.text())
            w4 = float(self.lineEdit_2.text())
        except:
            w1 = self.cost["lateral"]["w1"]
            w2 = self.cost["lateral"]["w2"]
            w3 = self.cost["lateral"]["w3"]
            w4 = self.cost["lateral"]["w4"]
        state_cost = [w1, w2, w3]
        control_cost = [w4]
        return state_cost, control_cost

    def get_longitudinal_cost(self):
        try:
            w1 = float(self.lineEdit_6.text())
            w2 = float(self.lineEdit_7.text())
            w3 = float(self.lineEdit_8.text())
            w4 = float(self.lineEdit_9.text())
        except:
            w1 = self.cost["longitudinal"]["w1"]
            w2 = self.cost["longitudinal"]["w2"]
            w3 = self.cost["longitudinal"]["w3"]
            w4 = self.cost["longitudinal"]["w4"]
        state_cost = [w1, w2, w3]
        control_cost = [w4]
        return state_cost, control_cost

    def calculate_transverse_mpc(self):
        time_begin = time.time()
        # 触发横向MPC计算
        self.transverse = ProcessedMatrix()
        self.transverse.init_process(N=self.N, vehicle_conf=self.vehicle_params, dt=self.dt)
        tran_state = np.array([[0],
                               [0],
                               [self.current_control_steer if self.current_control_steer else 0]])

        # tran_state = np.array([[float(self.processed_measurement.get_predict_traj()['y'][0])],
        #                        [float(self.processed_measurement.get_predict_traj()['h'][0])],
        #                        [self.current_control_steer if self.current_control_steer else 0]])
        tran_control = np.array([[self.current_control_steer if self.current_control_steer else 0]])
        self.transverse.init_process(tran_state, tran_control, v_ref=list(self.current_predict_traj['v']))
        state_cost, control_cost = self.get_transverse_cost()
        self.transverse.calculate_qn(state_cost)
        self.transverse.calculate_rn(control_cost)
        self.transverse.calculate_x_ref(self.processed_measurement.get_predict_traj())
        self.transverse.calculate_H()
        self.transverse.calculate_g()
        self.transverse.calculate_Acon()
        self.transverse.calculate_min_max(- self.vehicle_params.vehicle_params["maximum_angle"],
                                          self.vehicle_params.vehicle_params["maximum_angle"],
                                          state_min=[-10, -30, -30], state_max=[10, 30, 30])
        self.transverse.update_qp()
        if self.transverse.res.info.status != "solved":
            self.label_23.setText("false!")
        else:
            time_end = time.time()
            time_span = str(round((time_end - time_begin) * 1000, 2)) + "ms"
            self.label_23.setText(time_span)
        self.disassembly_condition(1)

    def calculate_longitudinal_mpc(self):
        # 触发纵向MPC计算
        begin_time = time.time()
        self.longitudinal = ProcessedMatrix()
        self.longitudinal.init_process(N=self.N, vehicle_conf=self.vehicle_params, dt=self.dt)
        lon_state = np.array([[0],
                               [float(self.current_state_mpc['v'][0])],
                               [self.current_control_steer if self.current_control_steer else 0]])
        lon_control = np.array([[self.current_control_steer if self.current_control_steer else 0]])
        self.longitudinal.init_process(xk=lon_state, uk=lon_control)
        state_cost, control_cost = self.get_longitudinal_cost()
        self.longitudinal.calculate_qn(state_cost)
        self.longitudinal.calculate_rn(control_cost)
        self.longitudinal.calculate_x_ref_lon(self.processed_measurement.get_predict_traj())
        self.longitudinal.calculate_H()
        self.longitudinal.calculate_g()
        self.longitudinal.calculate_Acon(lat=False)
        self.longitudinal.calculate_min_max(- self.vehicle_params.vehicle_params["maximum_acceleration"],
                                            self.vehicle_params.vehicle_params["maximum_acceleration"],
                                            state_min=
                                            [0, 0, -self.vehicle_params.vehicle_params["maximum_acceleration"]],
                                            state_max=[1000, self.vehicle_params.vehicle_params["maximum_velocity"],
                                                       self.vehicle_params.vehicle_params["maximum_acceleration"]])
        self.longitudinal.update_qp()
        if self.longitudinal.res.info.status != "solved":
            self.label_6.setText("false!")
        else:
            time_end = time.time()
            time_span = str(round((time_end - begin_time) * 1000, 2)) + "ms"
            self.label_6.setText(time_span)
        self.disassembly_condition(2)

    def generate_predict_traj(self):
        y, h = self.disassembly_condition(1)
        s, v = self.disassembly_condition(2)
        y_list = [0]
        x_list = [0]
        for i in range(1, self.N):
            y_list.append(y[i] + y_list[i - 1])
            z = math.sqrt(abs((s[i] - s[i - 1]) ** 2 - y[i] ** 2))
            x_list.append(x_list[i - 1] + z)
        new_x = []
        new_y = []
        for i in range(self.N):
            new_x.append(x_list[i] * np.cos(self.processed_measurement.match_state['h'] + np.pi) + y_list[i] * np.sin(
                self.processed_measurement.match_state['h']+ np.pi))
            new_y.append(x_list[i] * np.sin(self.processed_measurement.match_state['h']+ np.pi) - y_list[i] * np.cos(
                self.processed_measurement.match_state['h']+ np.pi))

        for i in range(self.N):
            new_x[i] += float(self.processed_measurement.match_state['x'])
            new_y[i] += float(self.processed_measurement.match_state['y'])
        if self.predict_traj is not None:
            self.pen.removeItem(self.predict_traj)
        self.predict_traj = pg.PlotCurveItem(new_x, new_y, pen=pg.mkPen(color=(110, 73, 110), width=2, symbolBrush='b'))
        self.pen.addItem(self.predict_traj)


    def disassembly_condition(self, mode):
        # 获取预测状态与实际状态
        if mode == 1 and self.transverse.predicted_state is not None and any(self.transverse.predicted_state):
            y = []
            h = []
            steer = []
            control_steer = []
            y_ref = []
            h_ref = []
            steer_ref = []
            for i in range(self.N):
                y_ref.append(
                    float(self.transverse.x_ref[i * self.transverse.shape: i * self.transverse.shape + 1, :][0, 0]))
                h_ref.append(
                    float(self.transverse.x_ref[i * self.transverse.shape + 1: i * self.transverse.shape + 2, :][0, 0]))
                steer_ref.append(
                    float(self.transverse.x_ref[i * self.transverse.shape + 2: i * self.transverse.shape + 3, :][0, 0]))
                y.append(float(self.transverse.predicted_state[i * self.transverse.shape][0]))
                h.append(float(self.transverse.predicted_state[i * self.transverse.shape + 1][0]))
                steer.append(float(self.transverse.predicted_state[i * self.transverse.shape + 2][0]))
                control_steer.append(float(self.transverse.predicted_control[i][0]))
            if self.lateral_windows is None:
                self.init_windows_lateral()
            self.lateral_windows.update_data_ref(y_ref, h_ref, steer_ref)
            self.lateral_windows.update_data_pre(y, h, control_steer)
            return y, h

        if mode == 2 and self.longitudinal.predicted_state is not None and any(self.longitudinal.predicted_state):
            x = []
            v = []
            a = []
            control_a = []
            x_ref = []
            v_ref = []
            a_ref = []
            for i in range(self.N):
                x_ref.append(
                    float(self.longitudinal.x_ref[i * self.transverse.shape: i * self.transverse.shape + 1, :][0, 0]))
                v_ref.append(
                    float(self.longitudinal.x_ref[i * self.transverse.shape + 1: i * self.transverse.shape + 2, :][0, 0]))
                a_ref.append(
                    float(self.longitudinal.x_ref[i * self.transverse.shape + 2: i * self.transverse.shape + 3, :][0, 0]))
                x.append(float(self.longitudinal.predicted_state[i * self.transverse.shape][0]))
                v.append(float(self.longitudinal.predicted_state[i * self.transverse.shape + 1][0]))
                a.append(float(self.longitudinal.predicted_state[i * self.transverse.shape + 2][0]))
                control_a.append(float(self.longitudinal.predicted_control[i][0]))
            if self.longitudinal_windows is None:
                self.init_windows_longitudinal()
            self.longitudinal_windows.update_data_ref(x_ref, v_ref, a_ref)
            self.longitudinal_windows.update_data_pre(x, v, a)
            return x, v

    def closeEvent(self, event: QEvent) -> None:
        if self.longitudinal_windows is not None:
            self.longitudinal_windows.close()
        if self.lateral_windows is not None:
            self.lateral_windows.close()
        super(Visual, self).closeEvent(event)


if __name__ == '__main__':
    app = QApplication(sys.argv)  # 创建一个应用程序对象
    ui = Visual()
    ui.setObjectName("地图")
    ui.show()  # 展示窗口
    sys.exit(app.exec_())
